import logging

from langchain.embeddings.base import Embeddings
from langchain.pydantic_v1 import BaseModel, root_validator
from typing import Any, Dict, List

logger = logging.getLogger(__name__)


class ZhipuAIEmbeddings(BaseModel, Embeddings):
    """`Zhipu Embeddings` embedding model"""

    client: Any
    """zhiupai.ZhipuAI"""

    @root_validator(allow_reuse=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """实例化ZhipuAI为values["client"]

        Arguments:
            values -- 包含配置信息的字典，必须包含client字段

        Returns:
            values -- 包含配置信息的字典，如果环境中有zhipuai库，则将返回实例化的ZhipuAI类；否则将报错 'ModuleNotFoundError: No module named 'zhipuai''.
        """
        from zhipuai import ZhipuAI

        values["client"] = ZhipuAI()
        return values

    def embed_query(self, text: str) -> List[float]:
        """生成输入文本的embedding.

        Arguments:
            text -- 要生成embedding的文本

        Returns:
            embeddings -- 输入文本的embedding，一个浮点数值列表
        """
        embeddings = self.client.embeddings.create(
            model="embedding-2",
            input=text,
        )
        return embeddings.data[0].embedding

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """生成文本列表的embedding

        Arguments:
            text -- 要生成embedding的文本列表

        Returns:

            输入列表中每个文档的embedding列表，一个浮点数值列表的列表
        """
        return [self.embed_query(text) for text in texts]

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """Asynchronous Embed search docs."""
        raise NotImplementedError(
            "Please use `embed_documents`. Official does not support asynchronous requests"
        )

    async def aembed_query(self, text: str) -> List[float]:
        """Asynchronous Embed query text."""
        raise NotImplementedError(
            "Please use `aembed_query`. Official does not support asynchronous requests"
        )
